import numpy as np
from scipy import signal
from typing import Union, List, Tuple, Optional
import logging
import sys
import os
from collections import deque

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
if project_root not in sys.path:
    sys.path.append(project_root)

from exception.base_exceptions import SignalProcessingException
from exception.exception_handler import global_exception_handler
from utils.FrequencyBandCalculator import FrequencyBandCalculator
from response.BrainWaveResult import BrainWaveResult

logger = logging.getLogger(__name__)

# 定义自定义异常
class FilterException(SignalProcessingException):
    pass

class NotchFilter:
    """
    陷波滤波器工具类
    
    专门用于去除脑电信号中的工频干扰（50Hz/60Hz），
    可选择性使用，与service层完全解耦。
    """
    
    def __init__(self, sampling_rate: float = 512.0, notch_freq: float = 50.0, 
                 quality_factor: float = 30.0, filter_order: int = 4):
        """
        初始化陷波滤波器
        
        Args:
            sampling_rate: 采样率 (Hz)
            notch_freq: 陷波频率 (Hz) - 50Hz(欧洲/中国) 或 60Hz(美国)
            quality_factor: 品质因子，控制陷波宽度，越大越窄
            filter_order: 滤波器阶数，影响滤波效果
        """
        self.sampling_rate = sampling_rate
        self.notch_freq = notch_freq
        self.quality_factor = quality_factor
        self.filter_order = filter_order
        
        # 设计滤波器系数
        self._design_filter()
        
        # 滤波器状态（用于实时滤波）
        self.zi = None
        self.is_initialized = False
        
        # 时域信号缓冲区（用于重新计算频段功率）
        self.raw_signal_buffer = deque(maxlen=int(sampling_rate))  # 保存1秒的数据用于FFT计算
        
        # 频段功率计算器
        self.band_calculator = FrequencyBandCalculator(sampling_rate=sampling_rate)
        
        logger.info(f"陷波滤波器初始化: {notch_freq}Hz, Q={quality_factor}, 采样率={sampling_rate}Hz")
    
    def _design_filter(self):
        """设计陷波滤波器"""
        try:
            # 计算归一化频率
            nyquist = self.sampling_rate / 2.0
            # 使用 `iirnotch` 更精确地设计陷波滤波器
            self.b, self.a = signal.iirnotch(
                self.notch_freq / nyquist,
                self.quality_factor,
            )
            
            logger.info(f"陷波滤波器设计完成: 陷波频率 {self.notch_freq}Hz")
            
        except Exception as e:
            error_msg = f"陷波滤波器设计失败，陷波频率: {self.notch_freq}Hz，原因: {str(e)}"
            filter_error = FilterException(error_msg, processing_step="_design_filter")
            global_exception_handler.handle_exception(filter_error)
            # 创建通过滤波器作为备选
            self.b, self.a = [1], [1]
    
    def filter_single_value(self, value: Union[int, float]) -> float:
        """
        实时滤波单个数据点
        
        Args:
            value: 输入信号值
            
        Returns:
            float: 滤波后的值
        """
        try:
            if not self.is_initialized:
                # 初始化滤波器状态
                self.zi = signal.lfilter_zi(self.b, self.a) * value
                self.is_initialized = True
            
            # 应用滤波器
            filtered_value, self.zi = signal.lfilter(
                self.b, self.a, [value], zi=self.zi
            )
            
            return float(filtered_value[0])
            
        except Exception as e:
            error_msg = f"实时滤波失败，输入值: {value}，原因: {str(e)}"
            filter_error = FilterException(error_msg, processing_step="filter_single_value")
            global_exception_handler.handle_exception(filter_error)
            return float(value)  # 返回原值
    
    def filter_signal_batch(self, signal_data: Union[List, np.ndarray]) -> np.ndarray:
        """
        批量滤波信号数据
        
        Args:
            signal_data: 输入信号数组
            
        Returns:
            np.ndarray: 滤波后的信号数组
        """
        try:
            signal_array = np.array(signal_data, dtype=float)
            
            if len(signal_array) == 0:
                return signal_array
            
            # 应用零相位滤波（前向后向滤波）
            filtered_signal = signal.filtfilt(self.b, self.a, signal_array)
            
            return filtered_signal
            
        except Exception as e:
            error_msg = f"批量滤波失败，信号长度: {len(signal_data) if hasattr(signal_data, '__len__') else 'unknown'}，原因: {str(e)}"
            filter_error = FilterException(error_msg, processing_step="filter_signal_batch")
            global_exception_handler.handle_exception(filter_error)
            return np.array(signal_data, dtype=float)
    
    def get_frequency_response(self, frequencies: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        获取滤波器频率响应
        
        Args:
            frequencies: 频率点数组，默认自动生成
            
        Returns:
            tuple: (频率数组, 幅度响应数组)
        """
        if frequencies is None:
            frequencies = np.logspace(-1, np.log10(self.sampling_rate/2), 1000)
        
        nyquist = self.sampling_rate / 2.0
        w, h = signal.freqz(self.b, self.a, worN=frequencies/nyquist * np.pi)
        
        return frequencies, np.abs(h)
    
    def reset_filter_state(self):
        """重置滤波器状态"""
        self.zi = None
        self.is_initialized = False
        self.raw_signal_buffer.clear()
        logger.info("陷波滤波器状态已重置")
    
    def update_notch_frequency(self, new_freq: float):
        """
        更新陷波频率并重新设计滤波器
        
        Args:
            new_freq: 新的陷波频率
        """
        self.notch_freq = new_freq
        self._design_filter()
        self.reset_filter_state()
        logger.info(f"陷波频率已更新为: {new_freq}Hz")
    
    def get_filter_info(self) -> dict:
        """获取滤波器信息"""
        return {
            'sampling_rate': self.sampling_rate,
            'notch_frequency': self.notch_freq,
            'quality_factor': self.quality_factor,
            'filter_order': self.filter_order,
            'bandwidth': self.notch_freq / self.quality_factor,
            'is_initialized': self.is_initialized
        }

    # 为了简洁，暂时移除 filter_multi_channel, filter_brain_wave_result, process_8bands
    # 你可以根据需要将其移回或在外部调用